#include "audio_processor.h"
#include "audio_sr.h"
#include "audio_decoder.h"
#include "audio_encoder.h"
#include "bsp_codec.h"
#include <stdlib.h>
#include <freertos/FreeRTOS.h>
#include <freertos/task.h>
#include <freertos/ringbuf.h>
#include <esp_heap_caps.h>
#include <string.h>
#include <esp_log.h>

#define TAG "audio_processor"

typedef struct
{
    audio_encoder_t *audio_encoder;
    audio_decoder_t *audio_decoder;
    audio_sr_t *audio_sr;

    RingbufHandle_t mic_buf;
    RingbufHandle_t enc_buf;
    RingbufHandle_t dec_buf;
    RingbufHandle_t speaker_buf;

} audio_processor_t;

static audio_processor_t *s_audio_processor = NULL;

static void audio_processor_speak_task(void *arg)
{
    void *buf = NULL;
    size_t len = 0;
    while (1)
    {
        buf = xRingbufferReceive(s_audio_processor->speaker_buf, &len, portMAX_DELAY);
        if (buf != NULL)
        {
            ESP_LOGD(TAG, "speak task, len %u", len);
            bsp_codec_write(buf, len);
            vRingbufferReturnItem(s_audio_processor->speaker_buf, buf);
        }
    }
}

void audio_processor_init(void)
{
    if (s_audio_processor)
    {
        return;
    }

    s_audio_processor = (audio_processor_t *)malloc(sizeof(audio_processor_t));
    assert(s_audio_processor);

    s_audio_processor->mic_buf = xRingbufferCreateWithCaps(1024 * 8, RINGBUF_TYPE_BYTEBUF, MALLOC_CAP_SPIRAM);
    assert(s_audio_processor->mic_buf);
    s_audio_processor->speaker_buf = xRingbufferCreateWithCaps(1024 * 8, RINGBUF_TYPE_NOSPLIT, MALLOC_CAP_SPIRAM);
    assert(s_audio_processor->speaker_buf);
    s_audio_processor->dec_buf = xRingbufferCreateWithCaps(1024 * 8, RINGBUF_TYPE_NOSPLIT, MALLOC_CAP_SPIRAM);
    assert(s_audio_processor->dec_buf);
    s_audio_processor->enc_buf = xRingbufferCreateWithCaps(1024 * 8, RINGBUF_TYPE_NOSPLIT, MALLOC_CAP_SPIRAM);
    assert(s_audio_processor->enc_buf);

    s_audio_processor->audio_sr = audio_sr_create();
    audio_sr_set_output_buffer(s_audio_processor->audio_sr, s_audio_processor->mic_buf);

    s_audio_processor->audio_encoder = audio_encoder_create();
    audio_encoder_set_input_buffer(s_audio_processor->audio_encoder, s_audio_processor->mic_buf);
    audio_encoder_set_output_buffer(s_audio_processor->audio_encoder, s_audio_processor->enc_buf);

    s_audio_processor->audio_decoder = audio_decoder_create();
    audio_decoder_set_input_buffer(s_audio_processor->audio_decoder, s_audio_processor->dec_buf);
    audio_decoder_set_output_buffer(s_audio_processor->audio_decoder, s_audio_processor->speaker_buf);
}

void audio_processor_start(void)
{
    assert(xTaskCreatePinnedToCoreWithCaps(audio_processor_speak_task, "spk_task", 1024 * 4, NULL, 5, NULL, 1, MALLOC_CAP_SPIRAM) == pdTRUE);
    audio_decoder_start(s_audio_processor->audio_decoder);
    audio_encoder_start(s_audio_processor->audio_encoder);
    audio_sr_start(s_audio_processor->audio_sr);
}

void *audio_processor_read(size_t *size_p)
{
    size_t rean_len = 0;
    void *read_buf = xRingbufferReceive(s_audio_processor->enc_buf, &rean_len, portMAX_DELAY);
    assert(read_buf);

    void *ret_buf = malloc(rean_len);
    assert(ret_buf);
    memcpy(ret_buf, read_buf, rean_len);
    *size_p = rean_len;
    vRingbufferReturnItem(s_audio_processor->enc_buf, read_buf);

    return ret_buf;
}

void audio_processor_write(void *buf, size_t size)
{
    assert(xRingbufferSend(s_audio_processor->dec_buf, buf, size, portMAX_DELAY) == pdTRUE);
}

void audio_processor_set_wakenet_callback(void (*callback)(void *), void *arg)
{
    assert(s_audio_processor);
    assert(s_audio_processor->audio_sr);
    audio_sr_set_wakenet_callback(s_audio_processor->audio_sr, callback, arg);
}

void audio_processor_set_vad_callback(void (*callback)(void *, vad_state_t vad_state), void *arg)
{
    assert(s_audio_processor);
    assert(s_audio_processor->audio_sr);
    audio_sr_set_vad_callback(s_audio_processor->audio_sr, callback, arg);
}

void audio_processor_reset_wakenet()
{
    assert(s_audio_processor);
    assert(s_audio_processor->audio_sr);
    audio_sr_reset_wakenet(s_audio_processor->audio_sr);
}
